#ifndef CELL_H
#define CELL_H

enum SplitMethod { MIDDLE, MEDIAN, MEAN };

#define INF 1.e200

#include <iostream>
#include <algorithm>
#include <complex>
#include <vector>
#include "boost/shared_ptr.hpp"

#include "dbg.h"
#include "Bounds.h"
#include "MemDebug.h"

const double PI = 3.141592653589793;
const double TWOPI = 2.*PI;
const double IOTA = 1.e-10;

// We use a code (to be used as a template parameter) to indicate which kind of data we
// are using for a particular use. 
// NData means just count the point.
// TData means use a scalar.  Nominally temperature, but works with any scalar.
// EData means use a shear. 
enum DataCode { NData=1 , TData=2 , EData=3 };


// This class encapsulates the differences in the different kinds of data being
// stored in a Cell.  It is used both for the input data from the file and also
// for the mean values for a given Cell.  Some extra useful information is sometimes
// also stored.
template <int DC, int M>
struct CellData;

template <int M>
struct CellData<NData,M>
{
    Position<M> pos;
    int n;

    CellData() {}

    CellData(const Position<M>& _pos) : pos(_pos), n(1) {}

    template <int M2>
    CellData(const Position<M2>& _pos) : pos(_pos), n(1) {}

    CellData(const std::vector<boost::shared_ptr<CellData<NData,M> > >& vdata,
             size_t start, size_t end) :
        n(end-start)
    {
        for(size_t i=start; i!=end; ++i) pos += vdata[i]->pos;
        pos /= n;
    }

};

template <int M>
struct CellData<TData,M> 
{
    Position<M> pos;
    double wt;
    double w;
    int n;

    CellData() {}

    CellData(const Position<M>& _pos, double _t, double _w) : 
        pos(_pos), wt(_t * _w), w(_w), n(1)
    {}

    template <int M2>
    CellData(const Position<M2>& _pos, double _t, double _w) : 
        pos(_pos), wt(_t * _w), w(_w), n(1)
    {}

    CellData(const std::vector<boost::shared_ptr<CellData<TData,M> > >& vdata,
             size_t start, size_t end) :
        pos(), wt(0.), w(0.), n(end-start)
    {
        for(size_t i=start;i<end;++i) {
            const CellData<TData,M>& data = *vdata[i];
            pos += data.pos * data.w;
            wt += data.wt;
            w += data.w;
        }
        pos /= w;
    }
};

template <int M>
struct CellData<EData,M> 
{
    Position<M> pos;
    std::complex<double> we;
    double w;
    double enorm;
    int n;

    CellData() {}

    CellData(const Position<M>& _pos, const std::complex<double>& _e, double _w) : 
        pos(_pos), we(_e * _w), w(_w), enorm(std::norm(_e)), n(1)
    {}

    template <int M2>
    CellData(const Position<M2>& _pos, const std::complex<double>& _e, double _w) : 
        pos(_pos), we(_e * _w), w(_w), enorm(std::norm(_e)), n(1)
    {}

    CellData(const std::vector<boost::shared_ptr<CellData<EData,M> > >& vdata,
             size_t start, size_t end) :
        pos(), we(0.), w(0.), n(end-start)
    {
        for(size_t i=start;i<end;++i) {
            const CellData<EData,M>& data = *vdata[i];
            pos += data.pos * data.w;
            we += data.we;
            w += data.w;
        }
        enorm = std::norm(we/w);
        pos /= w;
    }
};


template <class CellType>
static int DoCountLeaves(const CellType* c)
{
    if (c->getLeft()) {
        Assert(c->getRight());
        return c->getLeft()->countLeaves() + c->getRight()->countLeaves();
    } else return 1;
}

template <class CellType>
static std::vector<const CellType*> DoGetAllLeaves(const CellType* c)
{
    std::vector<const CellType*> ret;
    if (c->getLeft().get()) {
        std::vector<const CellType*> temp = c->getLeft()->getAllLeaves();
        ret.insert(ret.end(),temp.begin(),temp.end());
        Assert(c->getRight());
        temp = c->getRight()->getAllLeaves();
        ret.insert(ret.end(),temp.begin(),temp.end()); 
    } else {
        Assert(!c->getRight().get());
        ret.push_back(c);
    }
    return ret;
}

template <int DC, int M>
struct DataCompare 
{
    int split;
    DataCompare(int s) : split(s) {}
    bool operator()(const boost::shared_ptr<CellData<DC,M> >& cd1,
                    const boost::shared_ptr<CellData<DC,M> >& cd2) const 
    { return cd1->pos.get(split) < cd2->pos.get(split); }
};

template <int DC, int M>
struct DataCompareToValue 
{
    int split;
    double splitvalue;

    DataCompareToValue(int s, double v) : split(s), splitvalue(v) {}
    bool operator()(const boost::shared_ptr<CellData<DC,M> >& cd) const 
    { return cd->pos.get(split) < splitvalue; }
};

template <int DC, int M>
inline size_t SplitCell(
    std::vector<boost::shared_ptr<CellData<DC,M> > >& vdata, SplitMethod sm, 
    size_t start, size_t end, const Position<M>& meanpos)
{
    size_t mid=0;

    Bounds<M> b;
    for(size_t i=start;i<end;++i) b += vdata[i]->pos;

    int split = b.getSplit();

    switch (sm) { // three different split methods
      case MIDDLE :
           { // Middle is the average of the min and max value of x or y
               double splitvalue = b.getMiddle(split);
               DataCompareToValue<DC,M> comp(split,splitvalue);
               typename std::vector<boost::shared_ptr<CellData<DC,M> > >::iterator middle = 
                   std::partition(vdata.begin()+start,vdata.begin()+end,comp);
               mid = middle - vdata.begin();
           } break;
      case MEDIAN :
           { // Median is the point which divides the group into equal numbers
               DataCompare<DC,M> comp(split);
               mid = (start+end)/2;
               typename std::vector<boost::shared_ptr<CellData<DC,M> > >::iterator middle = 
                   vdata.begin()+mid;
               std::nth_element(
                   vdata.begin()+start,middle,vdata.begin()+end,comp);
           } break;
      case MEAN :
           { // Mean is the weighted average value of x or y
               double splitvalue = meanpos.get(split);
               DataCompareToValue<DC,M> comp(split,splitvalue);
               typename std::vector<boost::shared_ptr<CellData<DC,M> > >::iterator middle = 
                   std::partition(vdata.begin()+start,vdata.begin()+end,comp);
               mid = middle - vdata.begin();
           } break;
      default :
           myerror("Invalid SplitMethod");
    }

    if (mid == start || mid == end) {
        // With duplicate entries, can get mid == start or mid == end. 
        // This should only happen if all entries in this set are equal.
        // So it should be safe to just take the mid = (start + end)/2.
        // But just to be safe, re-call this function with sm = MEDIAN to 
        // make sure.
        Assert(sm != MEDIAN);
        return SplitCell(vdata,MEDIAN,start,end,meanpos);
    }
    Assert(mid > start);
    Assert(mid < end);
    return mid;
}

template <int DC, int M>
class Cell
{
public:

    // A Cell contains the accumulated data for a bunch of galaxies.
    // It is characterized primarily by a centroid and a size.
    // The centroid is simply the weighted centroid of all the galaxy positions.
    // The size is the maximum deviation of any one of these galaxies 
    // from the centroid.  That is, all galaxies fall within a radius
    // size from the centroid.
    // The structure also keeps track of some averages and sums about
    // the galaxies which are used in the correlation function calculations.

    Cell(std::vector<boost::shared_ptr<CellData<DC,M> > >& vdata, 
         SplitMethod sm=MEAN, size_t start=0, size_t end=0);
    ~Cell() {}

    const CellData<DC,M>* getData() const { return _data.get(); }
    double getSize() const { return _size; }
    double getSizeSq() const { return _sizesq; }
    // For PairCells, getAllSize is different from getSize.
    double getAllSize() const { return _size; }

    const Cell<DC,M>* getLeft() const { return _left.get(); }
    const Cell<DC,M>* getRight() const { return _right.get(); }

    int countLeaves() const { return DoCountLeaves(this); }
    std::vector<const Cell<DC,M>*> getAllLeaves() const { return DoGetAllLeaves(this); }

protected:

    double _size;
    double _sizesq;

    boost::shared_ptr<CellData<DC,M> > _data;
    boost::shared_ptr<Cell<DC,M> > _left;
    boost::shared_ptr<Cell<DC,M> > _right;
};

template <int DC, int M>
inline Cell<DC,M>::Cell(std::vector<boost::shared_ptr<CellData<DC,M> > >& vdata,
                        SplitMethod sm, size_t start, size_t end) :
    _size(0.), _sizesq(0.)
{
    if (end == 0) end = vdata.size();
    Assert(vdata.size()>0);
    Assert(end <= vdata.size());
    Assert(end > start);

    if (end - start == 1) {
        _data = vdata[start];
    } else {
        _data.reset(new CellData<DC,M>(vdata,start,end));

        _sizesq = 0.;
        for(size_t i=start;i<end;++i) {
            double devsq = DistSq(_data->pos,vdata[i]->pos);
            if (devsq > _sizesq) _sizesq = devsq;
        }
        _size = sqrt(_sizesq);

        if (_size > 0.) {
            size_t mid = SplitCell(vdata,sm,start,end,_data->pos);
            try {
                _left.reset(new Cell<DC,M>(vdata,sm,start,mid));
                _right.reset(new Cell<DC,M>(vdata,sm,mid,end));
            } catch (std::bad_alloc) {
                myerror("out of memory - cannot create new Cell");
            }
        }
    }
}

#endif
